/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.seatunnel.connectors.seatunnel.elasticsearch.source;

import org.apache.seatunnel.api.configuration.ReadonlyConfig;
import org.apache.seatunnel.api.source.Collector;
import org.apache.seatunnel.api.source.SourceReader;
import org.apache.seatunnel.api.table.type.SeaTunnelRow;
import org.apache.seatunnel.api.table.type.SeaTunnelRowType;
import org.apache.seatunnel.connectors.seatunnel.elasticsearch.client.EsRestClient;
import org.apache.seatunnel.connectors.seatunnel.elasticsearch.config.ElasticsearchConfig;
import org.apache.seatunnel.connectors.seatunnel.elasticsearch.config.SearchApiTypeEnum;
import org.apache.seatunnel.connectors.seatunnel.elasticsearch.config.SearchTypeEnum;
import org.apache.seatunnel.connectors.seatunnel.elasticsearch.dto.source.PointInTimeResult;
import org.apache.seatunnel.connectors.seatunnel.elasticsearch.dto.source.ScrollResult;
import org.apache.seatunnel.connectors.seatunnel.elasticsearch.serialize.source.DefaultSeaTunnelRowDeserializer;
import org.apache.seatunnel.connectors.seatunnel.elasticsearch.serialize.source.ElasticsearchRecord;
import org.apache.seatunnel.connectors.seatunnel.elasticsearch.serialize.source.SeaTunnelRowDeserializer;

import org.apache.commons.lang3.StringUtils;

import lombok.extern.slf4j.Slf4j;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Deque;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;

@Slf4j
public class ElasticsearchSourceReader
        implements SourceReader<SeaTunnelRow, ElasticsearchSourceSplit> {

    SourceReader.Context context;

    private final ReadonlyConfig connConfig;

    private EsRestClient esRestClient;

    Deque<ElasticsearchSourceSplit> splits = new LinkedList<>();
    boolean noMoreSplit;

    private final long pollNextWaitTime = 1000L;

    public ElasticsearchSourceReader(SourceReader.Context context, ReadonlyConfig connConfig) {
        this.context = context;
        this.connConfig = connConfig;
    }

    @Override
    public void open() {
        esRestClient = EsRestClient.createInstance(this.connConfig);
    }

    @Override
    public void close() throws IOException {
        esRestClient.close();
    }

    @Override
    public void pollNext(Collector<SeaTunnelRow> output) throws Exception {
        synchronized (output.getCheckpointLock()) {
            ElasticsearchSourceSplit split = splits.poll();
            if (split != null) {
                SeaTunnelRowType seaTunnelRowType = split.getSeaTunnelRowType();
                ElasticsearchConfig sourceIndexInfo = split.getElasticsearchConfig();
                scrollSearchResult(seaTunnelRowType, sourceIndexInfo, output);
            } else if (noMoreSplit) {
                // signal to the source that we have reached the end of the data.
                log.info("Closed the bounded ELasticsearch source");
                context.signalNoMoreElement();
            } else {
                Thread.sleep(pollNextWaitTime);
            }
        }
    }

    private void scrollSearchResult(
            SeaTunnelRowType seaTunnelRowType,
            ElasticsearchConfig sourceIndexInfo,
            Collector<SeaTunnelRow> output) {

        SeaTunnelRowDeserializer deserializer =
                new DefaultSeaTunnelRowDeserializer(seaTunnelRowType);

        // SQL client
        if (SearchTypeEnum.SQL.equals(sourceIndexInfo.getSearchType())) {
            log.info("Using SQL query for index: {}", sourceIndexInfo.getIndex());
            ScrollResult scrollResult =
                    esRestClient.searchBySql(
                            sourceIndexInfo.getSqlQuery(), sourceIndexInfo.getScrollSize());

            outputFromScrollResult(scrollResult, sourceIndexInfo, output, deserializer);
            while (StringUtils.isNotEmpty(scrollResult.getScrollId())) {
                scrollResult =
                        esRestClient.searchWithSql(
                                scrollResult.getScrollId(), scrollResult.getColumnNodes());
                outputFromScrollResult(scrollResult, sourceIndexInfo, output, deserializer);
            }
        } else {
            // Check if we should use PIT API
            if (SearchApiTypeEnum.PIT.equals(sourceIndexInfo.getSearchApiType())) {
                log.info("Using Point-in-Time (PIT) API for index: {}", sourceIndexInfo.getIndex());
                searchWithPointInTime(sourceIndexInfo, output, deserializer);
            } else {
                log.info("Using Scroll API for index: {}", sourceIndexInfo.getIndex());
                ScrollResult scrollResult =
                        esRestClient.searchByScroll(
                                sourceIndexInfo.getIndex(),
                                sourceIndexInfo.getSource(),
                                sourceIndexInfo.getQuery(),
                                sourceIndexInfo.getScrollTime(),
                                sourceIndexInfo.getScrollSize());
                outputFromScrollResult(scrollResult, sourceIndexInfo, output, deserializer);
                while (scrollResult.getDocs() != null && !scrollResult.getDocs().isEmpty()) {
                    scrollResult =
                            esRestClient.searchWithScrollId(
                                    scrollResult.getScrollId(), sourceIndexInfo.getScrollTime());
                    outputFromScrollResult(scrollResult, sourceIndexInfo, output, deserializer);
                }
            }
        }
    }

    /**
     * Search using Point-in-Time API.
     *
     * @param sourceIndexInfo The Elasticsearch configuration
     * @param output The collector to output rows
     * @param deserializer The deserializer to convert Elasticsearch records to SeaTunnel rows
     */
    private void searchWithPointInTime(
            ElasticsearchConfig sourceIndexInfo,
            Collector<SeaTunnelRow> output,
            SeaTunnelRowDeserializer deserializer) {

        // Create a PIT
        String pitId =
                esRestClient.createPointInTime(
                        sourceIndexInfo.getIndex(), sourceIndexInfo.getPitKeepAlive());
        sourceIndexInfo.setPitId(pitId);
        log.info(
                "Created Point-in-Time with ID: {} for index: {}",
                pitId,
                sourceIndexInfo.getIndex());

        try {
            // Initial search
            PointInTimeResult pitResult =
                    esRestClient.searchWithPointInTime(
                            pitId,
                            sourceIndexInfo.getSource(),
                            sourceIndexInfo.getQuery(),
                            sourceIndexInfo.getPitBatchSize(),
                            null, // No search_after for first request
                            sourceIndexInfo.getPitKeepAlive());

            // Output the results
            outputFromPitResult(pitResult, sourceIndexInfo, output, deserializer);

            // Continue searching while there are more results
            while (pitResult.isHasMore()) {
                // Update the PIT ID and search_after values for the next request
                sourceIndexInfo.setPitId(pitResult.getPitId());
                sourceIndexInfo.setSearchAfter(pitResult.getSearchAfter());

                // Execute the next search
                pitResult =
                        esRestClient.searchWithPointInTime(
                                sourceIndexInfo.getPitId(),
                                sourceIndexInfo.getSource(),
                                sourceIndexInfo.getQuery(),
                                sourceIndexInfo.getPitBatchSize(),
                                sourceIndexInfo.getSearchAfter(),
                                sourceIndexInfo.getPitKeepAlive());

                // Output the results
                outputFromPitResult(pitResult, sourceIndexInfo, output, deserializer);
            }
        } finally {
            // Always clean up the PIT when done
            if (pitId != null) {
                try {
                    esRestClient.deletePointInTime(pitId);
                } catch (Exception e) {
                    log.warn("Failed to delete Point-in-Time with ID: " + pitId, e);
                }
            }
        }
    }

    private void outputFromScrollResult(
            ScrollResult scrollResult,
            ElasticsearchConfig elasticsearchConfig,
            Collector<SeaTunnelRow> output,
            SeaTunnelRowDeserializer deserializer) {
        List<String> source = elasticsearchConfig.getSource();
        String tableId = elasticsearchConfig.getCatalogTable().getTablePath().toString();
        for (Map<String, Object> doc : scrollResult.getDocs()) {
            SeaTunnelRow seaTunnelRow =
                    deserializer.deserialize(new ElasticsearchRecord(doc, source, tableId));
            output.collect(seaTunnelRow);
        }
    }

    /**
     * Output rows from a Point-in-Time search result.
     *
     * @param pitResult The Point-in-Time search result
     * @param elasticsearchConfig The Elasticsearch configuration
     * @param output The collector to output rows
     * @param deserializer The deserializer to convert Elasticsearch records to SeaTunnel rows
     */
    private void outputFromPitResult(
            PointInTimeResult pitResult,
            ElasticsearchConfig elasticsearchConfig,
            Collector<SeaTunnelRow> output,
            SeaTunnelRowDeserializer deserializer) {
        List<String> source = elasticsearchConfig.getSource();
        String tableId = elasticsearchConfig.getCatalogTable().getTablePath().toString();
        for (Map<String, Object> doc : pitResult.getDocs()) {
            SeaTunnelRow seaTunnelRow =
                    deserializer.deserialize(new ElasticsearchRecord(doc, source, tableId));
            output.collect(seaTunnelRow);
        }
    }

    @Override
    public List<ElasticsearchSourceSplit> snapshotState(long checkpointId) throws Exception {
        return new ArrayList<>(splits);
    }

    @Override
    public void addSplits(List<ElasticsearchSourceSplit> splits) {
        this.splits.addAll(splits);
    }

    @Override
    public void handleNoMoreSplits() {
        noMoreSplit = true;
    }

    @Override
    public void notifyCheckpointComplete(long checkpointId) throws Exception {}
}
